# -*- encoding: utf-8 -*-
"""
@File    : scale_shift.py
@Author  : lilong
@Time    : 2023/2/28 10:46 上午
"""

import numpy as np

import tensorflow as tf
from tensorflow.keras.layers import Dense, Layer
from tensorflow.keras import backend as K
from tensorflow.keras.models import Model
from keras_layer_normalization import LayerNormalization


class ScaleShift(Layer):
    """缩放平移变换层（Scale and shift）"""

    def __init__(self, **kwargs):
        super(ScaleShift, self).__init__(**kwargs)

    def build(self, input_shape):
        kernel_shape = (1,) * (len(input_shape) - 1) + (input_shape[-1],)
        self.log_scale = self.add_weight(name='log_scale',
                                         shape=kernel_shape,
                                         initializer='zeros')
        self.shift = self.add_weight(name='shift',
                                     shape=kernel_shape,
                                     initializer='zeros')

    def call(self, inputs, **kwargs):
        """为数据进行tf.exp的目的是既不改变数据之间的关系又可以压缩数据尺寸，消除异方差的以及转化计算方法
        """
        x_outs = tf.exp(self.log_scale) * inputs + self.shift
        return x_outs


# ----测试----
x_sample = np.array(
    [
        [[0, 0, 0, 1, 0, 1, 0, 0]],
        [[0, 0, 1, 0, 0, 0, 0, 0]]
    ])
x_sample = K.cast(x_sample, 'float32')
x = ScaleShift()(x_sample)
print(x.numpy())
# [[[0. 0. 0. 1. 0. 1. 0. 0.]]
#  [[0. 0. 1. 0. 0. 0. 0. 0.]]]


